{
 "cells": [
  {
   "cell_type": "markdown",
   "id": "3548e26e-bc2f-4efe-87eb-8b9e1d24e9fc",
   "metadata": {},
   "source": [
    "# Sparks\n",
    "A. Ghafarollahi, M.J. Buehler*\n",
    "\n",
    "Massachusetts Institute of Technology\n",
    "\n",
    "*mbuehler@MIT.EDU"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "47ac9218-4b71-49d9-a36d-80f9bae5e161",
   "metadata": {},
   "source": [
    "# OpenAI API Key"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 1,
   "id": "e9ba4e2d-41a0-4d54-8f99-fe2702bbef13",
   "metadata": {},
   "outputs": [],
   "source": [
    "import os\n",
    "import json\n",
    "\n",
    "with open('config.json') as f:\n",
    "    config = json.load(f)\n",
    "\n",
    "your_openai_api_key = config[\"your_openai_api_key\"]\n",
    "os.environ['OPENAI_API_KEY'] = your_openai_api_key\n",
    "\n",
    "chroma_key = config[\"your_chroma_key\"]\n",
    "os.environ['CHROMA_KEY'] = chroma_key"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "1b41d7d4-f561-4659-9559-a4098428a0e7",
   "metadata": {},
   "source": [
    "# Define Query, Tools, Constraints"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "id": "b817dd29-3434-4b85-af6d-2749edef3541",
   "metadata": {},
   "outputs": [],
   "source": [
    "import json\n",
    "import sys\n",
    "from pathlib import Path\n",
    "import shutil\n",
    "\n",
    "####################################################################################\n",
    "# the project directory where all files will be stored\n",
    "base_dir = 'protein_principle' \n",
    "####################################################################################\n",
    "# User query\n",
    "query = 'Discover a novel principle in short protein peptides.'\n",
    "####################################################################################\n",
    "# Build all tools\n",
    "from Sparks.Sparks_utils import Toolset, Tool, make_parameters\n",
    "\n",
    "toolset = Toolset()\n",
    "\n",
    "tool_1 = Tool(\n",
    "    name=\"calculate_force_from_seq\",\n",
    "    description=(\n",
    "        \"This function takes a protein sequence in one-letter amino acid code, \"\n",
    "        \"performs the required analysis or simulation, and returns the \"\n",
    "        \"non-dimensional protein force (defined as the maximum force along the protein unfolding force-extension curve) as a float. \"\n",
    "        \"The function might lead to None results. So, the output should always be checked.\"\n",
    "    ),\n",
    "    parameters=make_parameters(\n",
    "        sequence=(\"string\", \"Protein sequence in standard one-letter amino acid codes (e.g., 'MKTFFV...').\")\n",
    "    )\n",
    ")\n",
    "\n",
    "tool_2 = Tool(\n",
    "    name=\"calculate_energy_from_seq\",\n",
    "    description=(\n",
    "        \"This function takes a protein sequence in one-letter amino acid code, \"\n",
    "        \"performs the required analysis or simulation, and returns the \"\n",
    "        \"non-dimensional protein energy (defined as the area under the protein unfolding force-extension curve) as a float. \"\n",
    "        \"The function might lead to None results. So, the output should always be checked.\"\n",
    "    ),\n",
    "    parameters=make_parameters(\n",
    "        sequence=(\"string\", \"Protein sequence in standard one-letter amino acid codes (e.g., 'MKTFFV...').\")\n",
    "    )\n",
    ")\n",
    "\n",
    "tool_3 = Tool(\n",
    "    name=\"design_protein_from_length\",\n",
    "    description=(\n",
    "        \"This function creates random protein sequences of a given length. \"\n",
    "        \"It returns the amino acid sequence (string format) of the generated protein.\"\n",
    "    ),\n",
    "    parameters=make_parameters(\n",
    "        length=(\"integer\", \"Desired number of amino acids in the protein sequence\")\n",
    "    )\n",
    ")\n",
    "\n",
    "tool_4 = Tool(\n",
    "    name=\"design_protein_from_CATH\",\n",
    "    description=(\n",
    "        \"This function creates random alpha-helix, beta-sheet, or mixed protein sequences of a given length. \"\n",
    "        \"It returns the amino acid sequence (string format) of the generated protein. \"\n",
    "        \"The function can be used to create proteins biased toward alpha-helix (CATH=1), beta-sheet (CATH=2), or mixed structures (CATH=3). \"\n",
    "        \"However, the output may not fully reflect the desired structural class, especially for short sequences. \"\n",
    "        \"To ensure the proteins follow the desired structure, the secondary structure of folded sequences should be analyzed.\"\n",
    "    ),\n",
    "    parameters=make_parameters(\n",
    "        length=(\"integer\", \"Desired number of amino acids in the protein sequence\"),\n",
    "        cath=(\"string\", 'Structure bias: \"1\" for alpha-helix, \"2\" for beta-sheet, \"3\" for mixed')\n",
    "    )\n",
    ")\n",
    "\n",
    "tool_5 = Tool(\n",
    "    name=\"analyze_protein_structure\",\n",
    "    description=(\n",
    "        \"This function returns the secondary structure percentage content of a protein from a PDB file. \"\n",
    "        \"The output is a dictionary with 8-class DSSP secondary structure elements as keys: \"\n",
    "        '{\"H\", \"B\", \"E\", \"G\", \"I\", \"T\", \"S\", \"P\", \"-\"}. Each key maps to the corresponding percentage value (float).'\n",
    "    ),\n",
    "    parameters=make_parameters(\n",
    "        protein_struct=(\"string\", \"Protein structure file in PDB format\")\n",
    "    )\n",
    ")\n",
    "\n",
    "tool_6 = Tool(\n",
    "    name=\"fold_protein\",\n",
    "    description=(\n",
    "        \"This function takes a protein amino acid sequence, folds it and creates the protein PDB file. \"\n",
    "        \"Returns the PDB file name (string format) of the resulting 3D folded structure. \"\n",
    "        \"The sequence should be in FASTA format and contain only valid amino acid characters (no special symbols like '-').\"\n",
    "    ),\n",
    "    parameters=make_parameters(\n",
    "        sequence=(\"string\", \"Protein sequence in FASTA format (only amino acid characters, no special symbols).\"),\n",
    "        name=(\"string\", \"Name to be used for saving the folded protein structure.\")\n",
    "    )\n",
    ")\n",
    "\n",
    "toolset.add(tool_1)\n",
    "toolset.add(tool_2)\n",
    "toolset.add(tool_3)\n",
    "toolset.add(tool_4)\n",
    "toolset.add(tool_5)\n",
    "toolset.add(tool_6)\n",
    "\n",
    "tools = toolset.as_json()\n",
    "####################################################################################\n",
    "# Computational constrainsts\n",
    "constraints = '''\n",
    "1- During the research, assume protein lengths to be equal or between 30 and 80.\n",
    "2- Consider protein sequences in intervals of 10.\n",
    "3- For computational efficiency, restrict the total number of samples per protein length to 10.\n",
    "'''\n",
    "constraints = json.dumps(constraints, indent=4)\n",
    "####################################################################################\n",
    "# Total number of follow-up rounds\n",
    "total_followup_rounds = 3\n",
    "####################################################################################"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "534a8815-2e29-40e5-b815-7d6a2fe5cc90",
   "metadata": {},
   "outputs": [],
   "source": [
    "os.environ[\"PROJECT_DIR\"] = str(Path(base_dir).resolve())\n",
    "\n",
    "try:\n",
    "    os.mkdir(base_dir)\n",
    "except:\n",
    "    print('Directory exists.')\n",
    "\n",
    "try:\n",
    "    shutil.copy('./Sparks/functions.py', base_dir)\n",
    "except:\n",
    "    pass\n",
    "\n",
    "try:\n",
    "    shutil.copy('./Sparks/MD_protein.py', base_dir)\n",
    "except:\n",
    "    pass\n",
    "\n",
    "try:\n",
    "    shutil.copy('./Sparks/utils_ForceGPT.py', base_dir)\n",
    "except:\n",
    "    pass\n",
    "\n",
    "shutil.copytree('./Sparks/0_codes', base_dir+'/0_codes', dirs_exist_ok=True)\n",
    "\n",
    "sys.path.append('./Sparks')\n",
    "from Sparks_functions import *"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "4f10e822-46ba-4cc8-b245-6b8d8602290a",
   "metadata": {},
   "outputs": [],
   "source": [
    "# Idea generation\n",
    "generate_idea(query, tools, constraints)\n",
    "\n",
    "# Idea testing\n",
    "test_idea(query, tools, constraints)\n",
    "initial_implementation()\n",
    "\n",
    "# Refinement\n",
    "last_exp_round = refine_approach(query, constraints, total_followup_rounds)\n",
    "\n",
    "# Create plots\n",
    "create_reason_plots(query, last_exp_round)\n",
    "\n",
    "# Documentation\n",
    "create_introduction(query)\n",
    "create_methods(query, tools, constraints, last_exp_round)\n",
    "create_results()\n",
    "create_conclusion(query)\n",
    "create_outlook(query, tools)\n",
    "create_pdf()"
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python 3 (ipykernel)",
   "language": "python",
   "name": "python3"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 3
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython3",
   "version": "3.10.13"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 5
}
